/*
 * @Description: 
 * @Author: Sang Hao
 * @Date: 2021-11-03 14:59:58
 * @LastEditTime: 2021-11-21 01:03:24
 * @LastEditors: Sang Hao
 */
#include "lidar_slam/models/segmentation/mesh_dbscan_segment.hpp"

namespace lidar_slam {

MeshDbscanSegment::MeshDbscanSegment(const YAML::Node& config_node) {
	min_sample_ = config_node["min_sample"].as<int>();
	radius_man_ = config_node["radius_man"].as<int>();

    threshold_angle_ = config_node["threshold_angle"].as<float>();
    threshold_z_ = config_node["threshold_z"].as<float>();
    threshold_d_ = config_node["threshold_d"].as<float>();

    threshold_pointnum_ = config_node["threshold_pointnum"].as<int>();
    
    float res_distance = config_node["res_distance"].as<float>();
    float res_angle = config_node["res_angle"].as<float>();

    polar_meshing_ptr_ = std::make_shared<PolarMeshing>(res_distance, res_angle);
}

MeshDbscanSegment::MeshDbscanSegment(int min_sample, int radius_man, float res_distance, float res_angle) 
	: min_sample_(min_sample), radius_man_(radius_man) {
    polar_meshing_ptr_ = std::make_shared<PolarMeshing>(res_distance, res_angle);
}


bool MeshDbscanSegment::GroundSeg(const CloudData::CLOUD_PTR& input_cloud_ptr, std::map<int, MeshData>& all_mesh_data) {
    bool distJudge, thetaJudge;
    for (auto it = all_mesh_data.begin(); it != all_mesh_data.end(); ++it) {
        if (input_cloud_ptr->points[it->second.zMinPtID].z < threshold_z_) { // 说明该网格内有地面点
            for (auto ite = it->second.PtID.begin(); ite != it->second.PtID.end(); ++ite) {
                distJudge = Dist3(input_cloud_ptr, *ite, it->second.zMinPtID) < threshold_d_;
                thetaJudge = Theta(input_cloud_ptr, *ite, it->second.zMinPtID) < threshold_angle_;
                if (distJudge || thetaJudge) { // 地面点判断条件
                    it->second.gPtID.insert(*ite);
                    // it->second.giMean += (input_cloud_ptr.points[*ite].intensity -
                        //  it->second.giMean) / it->second.gPtID.size();
                } else {
                    it->second.ngPtID.insert(*ite);
                    // it->second.ngiMean += (input_cloud_ptr.points[*ite].intensity - 
                        // it->second.ngiMean) / it->second.ngPtID.size();
                }
            }
        } else { // 说明该网格内都不是地面点
            for (auto ite = it->second.PtID.begin(); ite != it->second.PtID.end(); ++ite) {
                it->second.ngPtID.insert(*ite);
                // it->second.ngiMean += input_cloud_ptr.points[*ite].intensity;
            }
            // it->second.ngiMean /= it->second.ngPtID.size(); // 先累加，再一个除法，能快一些
        }
    }

    return true;
}


bool MeshDbscanSegment::NoGroundSeg(const CloudData::CLOUD_PTR& input_cloud_ptr, 
    std::map<int, MeshData>& all_mesh_data) {
    
    DbscanSeg(all_mesh_data);

    return true; 
}



void MeshDbscanSegment::GetOccuSet(std::map<int, MeshData>& all_mesh_data, std::set<int>& occu_set) {
    for (auto it = all_mesh_data.begin(); it != all_mesh_data.end(); ++it) {
        if (int(it->second.ngPtID.size()) > threshold_pointnum_) {
            occu_set.insert(it->first);
        }
    }
}

/* 这里暂时没有用曼哈顿距离，暂时简单的使用相邻 */
int MeshDbscanSegment::FindNeighbor(int query, std::set<int>& occu_set, std::set<int>& neighbor) {
    int mesh_width = polar_meshing_ptr_->Width();

    int row_i = query / mesh_width;
    int col_j = query % mesh_width;
    
    int n = radius_man_;
    int ite = 0;
    int num = 0;
    
    for (int ii = row_i - n; ii <= row_i + n; ++ii) {
        if (ii < 0) {
            continue;
        }
        ite = 0;
        for (int jj = col_j - n; ++ite <= (2 * n + 1); ++jj) {
            if (jj < 0)
                jj += mesh_width;
            if (jj >= mesh_width)
                jj -= mesh_width;
            if (occu_set.find(ii * mesh_width + jj) != occu_set.end()) {
                if (neighbor.find(ii * mesh_width + jj) == neighbor.end()) {
                    neighbor.insert(ii * mesh_width + jj);
                    ++num;
                }
                    
            }
        }
    }
    return num;
}


void MeshDbscanSegment::ClusterOne(int query, std::set<int>& occu_set, std::set<int>& temp, std::set<int>& clusterk) {
    int Ptnew;
    Ptnew = FindNeighbor(query, occu_set, temp);  // 先放入temp中
    // if (Ptnew == 0) return ++cluster;
    int next;
    while (1)
    {
        if (Ptnew == 0) break;
        auto it = temp.begin();
        next = *it;
        if (next == 0) break;
        temp.erase(next);   // 这里每删一个，cluout就要insert一个
        clusterk.insert(next);
        occu_set.erase(next); 
        ClusterOne(next, occu_set, temp, clusterk);
    }
}

// cMapOut:first参数为gridID
void MeshDbscanSegment::DbscanSeg(std::map<int, MeshData>& all_mesh_data) {
    std::set<int> occu_set;
    GetOccuSet(all_mesh_data, occu_set);
	int clusterID = 0;
	std::set<int> temp;
	std::set<int> cluster;
	while(occu_set.size() > 1) {
        auto ite = occu_set.begin();
        srand(time(0));
        advance(ite, rand() % occu_set.size());
        ClusterOne(*ite, occu_set, temp, cluster);
        clusterID++;
        for (auto it = cluster.begin(); it != cluster.end(); ++it) {
            all_mesh_data.find(*it)->second.label = clusterID;
        }
        temp.clear();
        cluster.clear();
    }
}

/**
 * @description: 
 * @param  {*}这个input_cloud_ptr总是没有分配内存
 * @return {*}
 */
bool MeshDbscanSegment::Segmentation(const CloudData::CLOUD_PTR& input_cloud_ptr,
		CloudRGBData::CLOUD_PTR& seged_cloud_ptr) {
    float  colors[] = {
		255, 0,   0,   // red 		0	
		0,   255, 0,   // green		1
		0,   0,   255, // blue		2
		255, 255, 0,   // yellow		3
		0,   255, 255, // light blue	4
		255, 0,   255, // magenta     5
		255, 128, 0,   // orange		6
		255, 153, 255, // pink		7
		51,  153, 255, //			8
		153, 102, 51,  //			9
		128, 51,  153, //			10
		153, 153, 51,  //			11
		163, 38,  51,  //			12
		204, 153, 102, //		13
		204, 224, 255, //		14
		128, 179, 255, //		15
		206, 255, 0,   //			16
		255, 204, 204, //			17
		204, 255, 153, //			18
		38,  128, 51,  // 		19

	}; // 20x3=60 color elements
	
    std::map<int, MeshData> all_mesh_data;
    polar_meshing_ptr_->CloudToMesh(input_cloud_ptr);
    polar_meshing_ptr_->GetAllMeshData(all_mesh_data);
    /* 先地面分割，此时all_mesh_data已经做了更新 */
    GroundSeg(input_cloud_ptr, all_mesh_data);

    DbscanSeg(all_mesh_data);

    for (auto it = all_mesh_data.begin(); it != all_mesh_data.end(); ++it) {
        /* 地面点 */
        for (auto itg = it->second.gPtID.begin(); itg != it->second.gPtID.end(); ++itg) {
            CloudRGBData::POINT p;
            p.x = input_cloud_ptr->points[*itg].x;
            p.y = input_cloud_ptr->points[*itg].y;
            p.z = input_cloud_ptr->points[*itg].z;
            /* 设置地面点为黑色 */
            // p.intensity = 255;
            p.r = 0;
            p.g = 0;
            p.b = 0;
            seged_cloud_ptr->points.push_back(p);
        }
        for (auto itng = it->second.ngPtID.begin(); itng != it->second.ngPtID.end(); ++itng) {
            if (it->second.label == 0) {
                continue;
            }
            CloudRGBData::POINT p;
            p.x = input_cloud_ptr->points[*itng].x;
            p.y = input_cloud_ptr->points[*itng].y;
            p.z = input_cloud_ptr->points[*itng].z;
            /* 设置非地面点颜色 */
            p.r = colors[it->second.label % 20 * 3 + 0];
            p.g = colors[it->second.label % 20 * 3 + 1];
            p.b = colors[it->second.label % 20 * 3 + 2];
            // p.intensity = intensitys[it->second.label % 5]; 
            seged_cloud_ptr->points.push_back(p);
        }
    }
    seged_cloud_ptr->height = 1;
	seged_cloud_ptr->width = seged_cloud_ptr->points.size();
    return true;        
}



/**
 * @description: 两个功能函数，
 * @param  {*}
 * @return {*}
 */

float MeshDbscanSegment::Dist3(const CloudData::CLOUD_PTR& input_cloud_ptr, int id1, int id2) {
	if (id1 == id2) return 0;
	float dx = input_cloud_ptr->points[id1].x - input_cloud_ptr->points[id2].x;
	float dy = input_cloud_ptr->points[id1].y - input_cloud_ptr->points[id2].y;
	float dz = input_cloud_ptr->points[id1].z - input_cloud_ptr->points[id2].z;
	return sqrt(dx * dx + dy * dy + dz * dz);
}

float MeshDbscanSegment::Theta(const CloudData::CLOUD_PTR& input_cloud_ptr, int id1, int id2) {
    if (id1 == id2) return 0;
	float dx = input_cloud_ptr->points[id1].x - input_cloud_ptr->points[id2].x;
	float dy = input_cloud_ptr->points[id1].y - input_cloud_ptr->points[id2].y;
	float dz = input_cloud_ptr->points[id1].z - input_cloud_ptr->points[id2].z;
	float theta = atan(fabs(dz) / sqrt(dx * dx + dy * dy)) * 180 / M_PI;
	return theta;
}


}
